package com.demo.shardingsphere5.service.impl;

import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import com.demo.shardingsphere5.dao.domain.User;
import com.demo.shardingsphere5.dao.mapper.UserMapper;
import com.demo.shardingsphere5.service.UserService;
import jakarta.annotation.Resource;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.reflection.SystemMetaObject;
import org.apache.shardingsphere.driver.jdbc.core.datasource.ShardingSphereDataSource;
import org.apache.shardingsphere.infra.database.DefaultDatabase;
import org.apache.shardingsphere.mode.manager.ContextManager;
import org.apache.shardingsphere.readwritesplitting.algorithm.loadbalance.RoundRobinReadQueryLoadBalanceAlgorithm;
import org.springframework.stereotype.Service;

import javax.sql.DataSource;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.Map;

/**
 * 静态读写分离
 * ShardingSphere5提供读写分离，基于轮询的读库负载均衡算法：{@link RoundRobinReadQueryLoadBalanceAlgorithm}
 * 更多负载均衡算法参考：https://shardingsphere.apache.org/document/current/cn/dev-manual/readwrite-splitting/
 *
 * 注意：使用读写分离时，在同一事物进行读写操作，所有操作都在主库进行
 */
@Service
public class ReadWriteStaticUserServiceImpl extends ServiceImpl<UserMapper, User>
        implements UserService {

    @Resource
    private DataSource dataSource;

    public void insert(Long userId) {
        User user = new User();
        user.setUserName("aaa");
        user.setUserId(userId);

        super.save(user);
    }

    public User querySlave(Long userId) {
        return super.lambdaQuery().eq(User::getUserId, userId).one();
    }

    public long queryMaster(Long userId) throws SQLException {
        String sql = "select count(*) from user where user_id = " + userId;

        try(ResultSet resultSet = execQuery("ds0", sql)) {
            while (resultSet.next()) {
                return resultSet.getLong(1);
            }
        }
        return 0;
    }

    public ResultSet execQuery(String logicDb, String sql) throws SQLException {
        DataSource dataSource = getDataSource(logicDb);

        return dataSource.getConnection().createStatement().executeQuery(sql);
    }

    public void execUpdate(String logicDb, String sql) throws SQLException {
        DataSource dataSource = getDataSource(logicDb);

        dataSource.getConnection().createStatement().executeUpdate(sql);
    }

    private DataSource getDataSource(String logicDb) {
        ShardingSphereDataSource shardingDataSource = (ShardingSphereDataSource) dataSource;
        MetaObject metaObject = SystemMetaObject.forObject(shardingDataSource);
        ContextManager contextManager = (ContextManager) metaObject.getValue("contextManager");
        // 没有配置spring.shardingsphere.database.name，使用默认值logic_db
        Map<String, DataSource> dataSourceMap = contextManager.getDataSourceMap(DefaultDatabase.LOGIC_NAME);
        return dataSourceMap.get(logicDb);
    }

}
